import numpy as np
import math
from scipy.interpolate import CubicSpline


class Spline:
    def __init__(
        self,
        points: list[tuple[float, float]],
        velocity: list[float] | float,
        distancePreStep: float,
    ) -> None:
        """
        points: list[tuple[float,float]], [[x1,y1],[x2,y2],...]
        """
        self.distancePreStep: float = distancePreStep
        self.x: list[float] = []
        self.y: list[float] = []
        self.v: list[float] = []
        self.theta: list[float] = []
        self.s: list[float] = []

        if isinstance(velocity, float):
            velocity = [velocity for i in range(len(points))]
        if len(points) < 2:
            raise ValueError("points must be at least 2 points")
        dis = [0.0]
        for i in range(1, len(points)):
            dis.append(
                dis[i - 1]
                + math.sqrt(
                    (points[i][0] - points[i - 1][0]) ** 2
                    + (points[i][1] - points[i - 1][1]) ** 2
                )
            )
        states = [
            (point[0], point[1], velocity) for point, velocity in zip(points, velocity)
        ]
        cs = CubicSpline(dis, states)
        self.length = dis[len(dis) - 1]
        samplerS = np.linspace(0, self.length, int(self.length / self.distancePreStep))
        states = [(s[0], s[1], s[2]) for s in cs(samplerS)]

        for i in range(len(states)):
            s = states[i]
            self.x.append(s[0])
            self.y.append(s[1])
            self.v.append(s[2])
            if i > 0:
                dx = states[i][0] - states[i - 1][0]
                dy = states[i][1] - states[i - 1][1]
                self.theta.append(
                    math.atan2(dy, dx) if s[2] >= 0 else math.atan2(-dy, -dx)
                )
                self.s.append(self.s[i - 1] + math.sqrt(dx**2 + dy**2))
            else:
                dx = states[1][0] - states[0][0]
                dy = states[1][1] - states[0][1]
                self.theta.append(
                    math.atan2(dy, dx) if s[2] >= 0 else math.atan2(-dy, -dx)
                )
                self.s.append(0)

    def drawPath(self, ax, color: str) -> None:
        ax.plot(self.x, self.y, color=color)


class ReferPath:
    def __init__(self, spline: Spline = None) -> None:
        self.spline = spline

    def update(self, spline: Spline) -> None:
        self.spline = spline

    def is_ready(self) -> bool:
        return self.spline is not None

    def sampler(self, distance: float) -> list[float]:
        """
        distance: float, distance in path
        return: list[float], [theta, v, x, y, s]
        """
        if self.spline is None:
            raise "spline is None"
        for i in range(len(self.spline.s)):
            if self.spline.s[i] > distance:
                return [
                    self.spline.theta[i],
                    self.spline.v[i],
                    self.spline.x[i],
                    self.spline.y[i],
                    self.spline.s[i],
                ]
        return [
            self.spline.theta[-1],
            self.spline.v[-1],
            self.spline.x[-1],
            self.spline.y[-1],
            self.spline.s[-1],
        ]

    def getNearestPoint(self, pos: tuple[float, float]) -> tuple[int, list[float]]:
        """
        pos: tuple[float,float], [x,y]
        return: tuple[int,list[float]], [index, [theta,v,x,y,s]]
        """
        if self.spline is None:
            raise "spline is None"
        dx = np.array(self.spline.x) - pos[0]
        dy = np.array(self.spline.y) - pos[1]
        d = dx**2 + dy**2
        index = np.argmin(d)
        return index.item(), [
            self.spline.theta[index],
            self.spline.v[index],
            self.spline.x[index],
            self.spline.y[index],
            self.spline.s[index],
        ]

    def getReferancePath(
        self, state: list[float], T: int, timeStep: float
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        return: tuple[np.ndarray,np.ndarray], [referStates,referControl]
        """
        if self.spline is None:
            raise "spline is None"
        index, start = self.getNearestPoint((state[2], state[3]))
        curS = start[4]  # s
        curV = start[1]  # use path velocity
        # curV = state[1] # use current velocity
        referStates = np.zeros((4, T + 1))
        referControl = np.zeros((2, T))
        for i in range(T + 1):
            states = self.sampler(
                curS + timeStep * abs(curV)
            )  # if v<0 but path should extend forward
            curS = states[4]
            curV = states[1]
            referStates[0, i] = states[0]
            referStates[1, i] = states[1]
            referStates[2, i] = states[2]
            referStates[3, i] = states[3]
        for i in range(T):
            if i < T - 1:
                referControl[0, i] = (
                    referStates[0, i + 1] - referStates[0, i]
                ) / timeStep
                referControl[1, i] = (
                    referStates[1, i + 1] - referStates[1, i]
                ) / timeStep
            else:
                referControl[:, i] = referControl[:, i - 1]
        return referStates, referControl

    def drawReferancePath(
        self,
        referState: np.ndarray,
        ax,
        color: str,
        index: bool = False,
        points: bool = False,
    ):
        import matplotlib.pyplot as plt

        x = referState[2, :]
        y = referState[3, :]
        ax.plot(x, y, color=color)
        if points:
            ax.scatter(x, y, color=color)
        if index:
            for i in range(len(x)):
                plt.text(x[i], y[i], str(i))


if __name__ == "__main__":
    from refer_path import ReferPath
    import matplotlib.pyplot as plt

    line = Spline([(0, 0), (1, 1), (2, 0), (3, -1), (4, 0)], 1, 0.02)
    path = ReferPath(line)
    start = (0.9, 0.5)
    T = 10
    referState, referControl = path.getReferancePath(start, T, 0.1)

    ax = plt.subplot(1, 2, 1)
    line.drawPath(ax, "blue")
    path.drawReferancePath(referState, ax, "green")
    ax.scatter(start[0], start[1], color="red")

    bx = plt.subplot(1, 2, 2)
    bx.plot(referState[0, :], label="theta")
    bx.plot(referState[1, :], label="v")
    bx.plot(referState[2, :], label="x")
    bx.plot(referState[3, :], label="y")
    bx.plot(referControl[0, :], label="w")
    bx.plot(referControl[1, :], label="a")
    bx.legend()

    plt.show()
